#!/bin/env python3
import ithildin as ith
import numpy as np
import matplotlib.pyplot as plt
import sys
from typing import List
from scipy.interpolate import RectBivariateSpline
from pydiffmap.diffusion_map import DiffusionMap
from pydiffmap import visualization as diff_visualization
import matplotlib

# AP model (phase_1000)

# Laadt eerst de data in
data = ith.SimData.from_stem("phase_1000/phase_1000")

# Vermijd randeffecten en initialiatie
def snij_randen_weg(data,variabele,begintijd):
    def begin(index):
        if index == 0:
            return begintijd
        return data.vars[variabele].shape[index]//8
    def eind(index):
        if index == 0:
            return data.vars[variabele].shape[index]
        if data.vars[variabele].shape[index] < 20:
            return data.vars[variabele].shape[index]
        return (7 * data.vars[variabele].shape[index])//8
    return data.vars[variabele][begin(0):eind(0),begin(1):eind(1),begin(2):eind(2),begin(3):eind(3)]


begintijd = 2
var_u = snij_randen_weg(data, 'u', begintijd).flatten()
var_v = snij_randen_weg(data, 'v', begintijd).flatten()

def als_kolomvector(waardes):
    return waardes.reshape((waardes.shape[0], 1))

# Verlaag aantal punten om geheugengebruik en tijdsduur te beperken.
# replace=True strikt genomen incorrect maar wel een stuk sneller
keuzes = np.random.choice(var_u.shape[0], 1000,replace=True)
var_u = var_u[keuzes]
var_v = var_v[keuzes]

# Plot het!
#matplotlib.use("TkAgg")
#plt.scatter(var_u,var_v)
#plt.show()

# Maak nu diffusion maps

faseruimte = np.column_stack((var_u,var_v))
del var_u, var_v # wat geheugen vrijgeven
dmap = DiffusionMap.from_sklearn(n_evecs=2)
dmap.fit(faseruimte[5:46,:])

# geen convergentie bij:
#   * 100
#   * 200
#   * 400
#
# wel convergentie bij:
#   * 4
#   * 5:46 (41 punten)

# (u,v)-scatterplot, gekleurd met eerste DC
diff_visualization.data_plot(dmap)

# Uitprobeersels:
#   * Door epsilon in te stellen op iets tussen 0.001 --- 10,
#     kan binnen redelijke tijd een antwoord gevonden worden
#     voor alle 1000 punten.
dmap = DiffusionMap.from_sklearn(epsilon=0.001)
dmap.fit(faseruimte)
diff_visualization.data_plot(dmap)
